#include "Parser.h"

namespace FlooP
{
	Parser::Parser (Lexer& lexer)
	:
		m_lexer (lexer)
	{
		// Fill up the buffers
		consume ();
		consume ();
		consume ();
		consume ();

		m_token_helpers[Token::NONE] = "NONE";
		m_token_helpers[Token::LEFT_PARENTHESES] = "(";
		m_token_helpers[Token::RIGHT_PARENTHESES] = ")";
		m_token_helpers[Token::LEFT_BRACKET] = "[";
		m_token_helpers[Token::RIGHT_BRACKET] = "]";
		m_token_helpers[Token::LEFT_BRACE] = "{";
		m_token_helpers[Token::RIGHT_BRACE] = "}";
		m_token_helpers[Token::COLON] = ":";
		m_token_helpers[Token::SEMICOLON] = ";";
		m_token_helpers[Token::COMMA] = ",";
		m_token_helpers[Token::DOT] = ".";
		m_token_helpers[Token::DEFINE] = "DEFINE";
		m_token_helpers[Token::PROCEDURE] = "PROCEDURE";
		m_token_helpers[Token::IF] = "IF";
		m_token_helpers[Token::THEN] = "THEN";
		m_token_helpers[Token::LOOP] = "LOOP";
		m_token_helpers[Token::MU_LOOP] = "MU-LOOP";
		m_token_helpers[Token::AT] = "AT";
		m_token_helpers[Token::MOST] = "MOST";
		m_token_helpers[Token::TIMES] = "TIMES";
		m_token_helpers[Token::BEGIN] = "BEGIN";
		m_token_helpers[Token::END] = "END";
		m_token_helpers[Token::BLOCK] = "BLOCK";
		m_token_helpers[Token::QUIT] = "QUIT";
		m_token_helpers[Token::ABORT] = "ABORT";
		m_token_helpers[Token::ARROW] = "<=";
		m_token_helpers[Token::EQUALS] = "=";
		m_token_helpers[Token::OR] = "OR";
		m_token_helpers[Token::AND] = "AND";
		m_token_helpers[Token::LESS_THAN] = "<";
		m_token_helpers[Token::GREATER_THAN] = ">";
		m_token_helpers[Token::PLUS] = "+";
		m_token_helpers[Token::STAR] = "*";
		m_token_helpers[Token::YES] = "YES";
		m_token_helpers[Token::NO] = "NO";
		m_token_helpers[Token::CELL] = "CELL";
		m_token_helpers[Token::OUTPUT] = "OUTPUT";
		m_token_helpers[Token::NUMBER] = "NUMBER";
		m_token_helpers[Token::STRING] = "STRING";
		m_token_helpers[Token::IDENTIFIER] = "IDENTIFIER";
		m_token_helpers[Token::END_OF_FILE] = "END OF FILE";
	}

	std::shared_ptr<AST> Parser::parse_program ()
	{
		std::vector<std::shared_ptr<ASTProcedure>> procedure;
		std::vector<std::shared_ptr<ASTProcedureCall>> calls;

		while (m_buffer[0].type () != Token::END_OF_FILE) {

			switch (m_buffer[0].type ())
			{
				case Token::DEFINE:
				{
					consume ();
					do_procedure (procedure);
				}
				break;

				default:
				{
					throw Error (m_lexer.line (), m_lexer.character (), "Parser : Procedure or call expected.");
				}
				break;
			}
		}

		return std::shared_ptr<AST> (new AST (procedure, calls));
	}

	void Parser::do_procedure (std::vector<std::shared_ptr<ASTProcedure>>& procedure)
	{
		require (Token::PROCEDURE);
		std::string name = require (Token::STRING).text ();
		require (Token::LEFT_BRACKET);

		std::vector<std::string> args;

		while (!match (Token::RIGHT_BRACKET)) {
			args.push_back (require (Token::IDENTIFIER).text ());

			if (match (Token::COMMA)) {
				consume ();

				if (!match (Token::IDENTIFIER)) {
					throw Error (m_lexer.line (), m_lexer.character (), "Parser : Argument name expected after colon.");
				}
			}
		}

		require (Token::RIGHT_BRACKET);
		require (Token::COLON);
		std::shared_ptr<ASTStatement> body = do_block ();
		require (Token::DOT);

		procedure.push_back (std::shared_ptr<ASTProcedure> (new ASTProcedure (name, args, body)));
	}

	std::shared_ptr<ASTStatement> Parser::do_stmt ()
	{
		std::shared_ptr<ASTStatement> stmt;

		// Not all statements require semicolons, which is weird, but thats how Hofstadter designed it
		switch (m_buffer[0].type ())
		{
			case Token::BLOCK  : stmt = do_block (); require (Token::SEMICOLON); break;
			case Token::IF	   : stmt = do_if (); break;
			case Token::LOOP   : stmt = do_loop (); break;
			case Token::MU_LOOP: stmt = do_muloop (); break;
			case Token::QUIT   : stmt = do_quit (); require (Token::SEMICOLON); break;
			case Token::ABORT  : stmt = do_abort (); require (Token::SEMICOLON); break;
			case Token::CELL   : stmt = do_assignment (); require (Token::SEMICOLON); break;
			case Token::OUTPUT : stmt = do_assignment (); require (Token::SEMICOLON); break;

			default:
			{
				throw Error (m_lexer.line (), m_lexer.character (), "Parser : Statement expected.");
			}
		}

		return stmt;
	}

	std::shared_ptr<ASTStatement> Parser::do_block ()
	{
		require (Token::BLOCK);
		uint32_t num = require (Token::NUMBER).number ();
		require (Token::COLON);
		require (Token::BEGIN);

		std::vector<std::shared_ptr<ASTStatement>> block;

		while (m_buffer[3].type () != Token::END) {
			if (match (Token::END_OF_FILE)) {
				throw Error (m_lexer.line (), m_lexer.character (), "Parser : 'block #: end' expected.");
			}

			block.push_back (do_stmt ());
		}

		require (Token::BLOCK);
		if (num != require (Token::NUMBER).number ()) {
			throw Error (m_lexer.line (), m_lexer.character (), "Parser : Block number %i does not have matching 'end' statement.", num);
		}

		require (Token::COLON);
		require (Token::END);

		return std::shared_ptr<ASTStatement> (new ASTBlock (block, num));
	}

	std::shared_ptr<ASTStatement> Parser::do_if ()
	{
		require (Token::IF);

		std::shared_ptr<ASTExpression> cond = do_expr ();

		require (Token::COMMA);
		require (Token::THEN);
		require (Token::COLON);

		std::shared_ptr<ASTStatement> if_part = do_stmt ();

		return std::shared_ptr<ASTStatement> (new ASTIf (cond, if_part));
	}

	std::shared_ptr<ASTStatement> Parser::do_loop ()
	{
		require (Token::LOOP);

		if (match (Token::AT)) {
			consume ();
			require (Token::MOST);
		}

		std::shared_ptr<ASTExpression> times = do_expr ();
		require (Token::TIMES);
		require (Token::COLON);
		std::shared_ptr<ASTStatement> body = do_stmt ();

		return std::shared_ptr<ASTStatement> (new ASTLoop (times, body));
	}

	std::shared_ptr<ASTStatement> Parser::do_muloop ()
	{
		require (Token::MU_LOOP);
		require (Token::COLON);
		std::shared_ptr<ASTStatement> body = do_stmt ();

		return std::shared_ptr<ASTStatement> (new ASTMuLoop (body));
	}

	std::shared_ptr<ASTStatement> Parser::do_quit ()
	{
		require (Token::QUIT);
		require (Token::BLOCK);

		uint32_t num = require (Token::NUMBER).number ();

		return std::shared_ptr<ASTStatement> (new ASTQuit (num));
	}

	std::shared_ptr<ASTStatement> Parser::do_abort ()
	{
		require (Token::ABORT);
		require (Token::LOOP);

		uint32_t num = require (Token::NUMBER).number ();

		return std::shared_ptr<ASTStatement> (new ASTAbort (num));
	}

	std::shared_ptr<ASTStatement> Parser::do_assignment ()
	{
		// For simplicity cell_num = 0 represents OUTPUT cell and cell_num = 1 represents CELL(0), etc.
		uint32_t cell_num = 0;

		// Check if we're assigning to a CELL
		if (match (Token::CELL)) {
			consume ();

			require (Token::LEFT_PARENTHESES);

			// Get the CELL number
			cell_num = require (Token::NUMBER).number () + 1;
			require (Token::RIGHT_PARENTHESES);
		} else {
			// If its not a cell we're assigning to, it must be OUTPUT
			require (Token::OUTPUT);
		}

		require (Token::ARROW);
		std::shared_ptr<ASTExpression> expr = do_expr ();

		return std::shared_ptr<ASTStatement> (new ASTAssignment (cell_num, expr));
	}

	// Handles <expression> "OR" <expression>
	std::shared_ptr<ASTExpression> Parser::do_expr ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr1 ();

		while (match (Token::OR)) {
			consume ();
			std::shared_ptr<ASTExpression> expr_right = do_expr1 ();

			expr = std::shared_ptr<ASTExpression> (new ASTCompare (ASTCompare::OR, expr, expr_right));
		}

		return expr;
	}

	// Handles <expression> "AND" <expression>
	std::shared_ptr<ASTExpression> Parser::do_expr1 ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr2 ();

		while (match (Token::AND)) {
			consume ();
			std::shared_ptr<ASTExpression> expr_right = do_expr2 ();

			expr = std::shared_ptr<ASTExpression> (new ASTCompare (ASTCompare::AND, expr, expr_right));
		}

		return expr;
	}

	// Handles <expression> "=" <expression>
	//         <expression> "<" <expression>
	//         <expression> ">" <expression>
	std::shared_ptr<ASTExpression> Parser::do_expr2 ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr3 ();

		switch (m_buffer[0].type ())
		{
			case Token::EQUALS:
			{
				consume ();
				std::shared_ptr<ASTExpression> expr_right = do_expr3 ();
				expr = std::shared_ptr<ASTExpression> (new ASTCompare (ASTCompare::EQUALS, expr, expr_right));
			}
			break;

			case Token::LESS_THAN:
			{
				consume ();
				std::shared_ptr<ASTExpression> expr_right = do_expr3 ();
				expr = std::shared_ptr<ASTExpression> (new ASTCompare (ASTCompare::LESS_THAN, expr, expr_right));
			}
			break;

			case Token::GREATER_THAN:
			{
				consume ();
				std::shared_ptr<ASTExpression> expr_right = do_expr3 ();
				expr = std::shared_ptr<ASTExpression> (new ASTCompare (ASTCompare::GREATER_THAN, expr, expr_right));
			}
			break;
		}

		return expr;
	}

	// Handles <expression> "+" <expression>
	std::shared_ptr<ASTExpression> Parser::do_expr3 ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr4 ();

		while (match (Token::PLUS)) {
			consume ();
			std::shared_ptr<ASTExpression> expr_right = do_expr4 ();
			expr = std::shared_ptr<ASTExpression> (new ASTBinaryMathOp (ASTBinaryMathOp::PLUS, expr, expr_right));
		}

		return expr;
	}

	// Handles <expression> "*" <expression>
	std::shared_ptr<ASTExpression> Parser::do_expr4 ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr5 ();

		while (match (Token::STAR)) {
			consume ();
			std::shared_ptr<ASTExpression> expr_right = do_expr5 ();
			expr = std::shared_ptr<ASTExpression> (new ASTBinaryMathOp (ASTBinaryMathOp::STAR, expr, expr_right));
		}

		return expr;
	}

	// Handles <identifier> "[" <argument_list> "]"
	std::shared_ptr<ASTExpression> Parser::do_expr5 ()
	{
		std::shared_ptr<ASTExpression> expr = do_expr6 ();

		if (match (Token::LEFT_BRACKET)) {
			consume ();
			if (expr->type () != ASTExpression::IDENTIFIER) {
				throw Error (m_lexer.line (), m_lexer.character (), "Parser : Procedure call expected.");
			}

			std::vector<std::shared_ptr<ASTExpression>> args;

			while (!match (Token::RIGHT_BRACKET)) {
				args.push_back (do_expr ());

				if (match (Token::COMMA)) {
					consume ();
				}
			}

			require (Token::RIGHT_BRACKET);

			expr = std::shared_ptr<ASTExpression> (new ASTProcedureCall (dynamic_cast<ASTIdentifier*>(expr.get ())->name (), args));
		}

		return expr;
	}

	// Handles <identifier>
	//		   <number>
	//		   <string>
	//		   "YES"
	//		   "NO"
	std::shared_ptr<ASTExpression> Parser::do_expr6 ()
	{
		switch (m_buffer[0].type ())
		{
			case Token::OUTPUT:
			{
				consume ();
				return std::shared_ptr<ASTExpression> (new ASTCell (0));
			}
			break;

			case Token::CELL:
			{
				consume ();

				require (Token::LEFT_PARENTHESES);
				std::shared_ptr<ASTExpression> expr = std::shared_ptr<ASTExpression> (new ASTCell (require (Token::NUMBER).number () + 1));
				require (Token::RIGHT_PARENTHESES);

				return expr;
			}
			break;

			case Token::IDENTIFIER:
			{
				std::string var = require (Token::IDENTIFIER).text ();
				return std::shared_ptr<ASTExpression> (new ASTIdentifier (var));
			}
			break;

			case Token::NUMBER:
			{
				uint32_t num = require (Token::NUMBER).number ();
				return std::shared_ptr<ASTExpression> (new ASTNumber (num));
			}
			break;

			case Token::STRING:
			{
				std::string str = require (Token::STRING).text ();
				return std::shared_ptr<ASTExpression> (new ASTString (str));
			}
			break;

			case Token::YES:
			{
				// Yes is represented as 1
				consume ();
				return std::shared_ptr<ASTExpression> (new ASTNumber (1));
			}
			break;

			case Token::NO:
			{
				// No is represented as 0
				consume ();
				return std::shared_ptr<ASTExpression> (new ASTNumber (0));
			}
			break;

			default:
			{
				throw Error (m_lexer.line (), m_lexer.character (), "Parser : Primary expression expected.");
			}
			break;
		}
	}

	bool Parser::match (Token::Type type)
	{
		return (m_buffer[0].type () == type);
	}

	Token Parser::require (Token::Type type)
	{
		if (m_buffer[0].type () != type) {
			throw Error (m_lexer.line (), m_lexer.character (), "Parser : '%s' was expected.", m_token_helpers[type].c_str ());
		} else {
			return consume ();
		}
	}

	Token Parser::consume ()
	{
		Token ret = m_buffer[0];

		m_buffer[0] = m_buffer[1];
		m_buffer[1] = m_buffer[2];
		m_buffer[2] = m_buffer[3];
		m_buffer[3] = m_lexer.read_token ();

		return ret;
	}
}